#ifdef USE_NETLINK
#include "nl.h"
#include <bits/sockaddr.h>
#include <linux/genetlink.h>
#include <linux/taskstats.h>
#include <sys/socket.h>
#include <errno.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <sys/syscall.h>
#include <string.h>
#include "utils.h"

#define GENLMSG_DATA(glh) ((void *)(NLMSG_DATA(glh) + GENL_HDRLEN))
#define GENLMSG_PAYLOAD(glh)  (NLMSG_PAYLOAD(glh, 0) - GENL_HDRLEN)
#define NLA_DATA(na)    ((void *)((char*)(na) + NLA_HDRLEN))
#define NLA_PAYLOAD(len)  (len - NLA_HDRLEN)
#define MAX_NL_MSG_SIZE 1024

struct nl_msg_template {
	struct nlmsghdr n;
	struct genlmsghdr g;
	char buf[MAX_NL_MSG_SIZE];
};

int send_nl_cmd(int sd, __u16 nlmsg_type, __u32 nlmsg_pid, __u8 genl_cmd, __u16 nla_type, 
				void *nla_data, int nla_len);

void setup_netlink(netlink_data_t* netlink_data, int pid, int bi) {
	struct sockaddr_nl local;
	CHECK_CALL_RET(socket(AF_NETLINK, SOCK_RAW, NETLINK_GENERIC), netlink_data->nl_sd);
	memset(&local, 0, sizeof(local));
	local.nl_family = AF_NETLINK;
	CHECK_CALL(bind(netlink_data->nl_sd, (struct sockaddr *) &local, sizeof(local)));
	// now get the netlink family id  for taskstats
	struct {
		struct nlmsghdr n;
		struct genlmsghdr g;
		char buf[256];
	} reply;
	char *family_name = TASKSTATS_GENL_NAME;
	send_nl_cmd(netlink_data->nl_sd, GENL_ID_CTRL, pid, CTRL_CMD_GETFAMILY, 
				CTRL_ATTR_FAMILY_NAME, (void *)family_name, strlen(TASKSTATS_GENL_NAME) + 1);
	int reply_len = recv(netlink_data->nl_sd, &reply, sizeof(reply), 0);
	if (reply.n.nlmsg_type == NLMSG_ERROR || (reply_len < 0) || !NLMSG_OK((&reply.n), reply_len)) {
		FAIL("balancer %d couldn't get netlink socket family id", bi);
	}
	struct nlattr *nattr = (struct nlattr *) GENLMSG_DATA(&reply);
	nattr = (struct nlattr *) ((char *) nattr + NLA_ALIGN(nattr->nla_len));
	if (nattr->nla_type == CTRL_ATTR_FAMILY_ID) netlink_data->nl_id = *(__u16 *) NLA_DATA(nattr);
	else {
		FAIL("cpu %d couldn't get netlink family id", bi);
	}
}

int get_task_stats_from_nl(netlink_data_t* netlink_data, pid_t pid, pid_t tid, 
						   double *utime, double *stime, double *rtime) {
	if (send_nl_cmd(netlink_data->nl_sd, netlink_data->nl_id, pid, 
					TASKSTATS_CMD_GET, TASKSTATS_CMD_ATTR_PID, &tid, sizeof(&tid)) < 0) {
		FAIL("PID %d getting nl data for task %d: %s", pid, tid, strerror(errno));
	}
	struct nl_msg_template msg;
	int reply_len = recv(netlink_data->nl_sd, &msg, sizeof(msg), 0);
	if (reply_len < 0) {
		WARN("netlink reply error: %d, error %s", errno, strerror(errno));
		return -errno;
	}
	if (msg.n.nlmsg_type == NLMSG_ERROR || !NLMSG_OK((&msg.n), reply_len)) {
		struct nlmsgerr *err = NLMSG_DATA(&msg);
		if (-err->error != 3 && -err->error != 1) {
			// this warning goes off when the thread being tracked dies at the end of a run, 
			// with error 3, or when operation is not permitted with error 1
			WARN("netlink reply error: %d, %s", -err->error, strerror(-err->error));
		}
		return err->error;
	}
	reply_len = GENLMSG_PAYLOAD(&msg.n);
	struct nlattr *nattr = (struct nlattr *)GENLMSG_DATA(&msg);
	int len = 0;
	while (len < reply_len) {
		len += NLA_ALIGN(nattr->nla_len);
		if (nattr->nla_type == TASKSTATS_TYPE_AGGR_PID) {
			int aggr_len = NLA_PAYLOAD(nattr->nla_len);
			int len2 = 0;
			nattr = (struct nlattr *)NLA_DATA(nattr);
			while (len2 < aggr_len) {
				if (nattr->nla_type == TASKSTATS_TYPE_STATS) {
					struct taskstats* tstats = (struct taskstats *)NLA_DATA(nattr);
					*utime = tstats->ac_utime * 1e-6;
					*stime = tstats->ac_stime * 1e-6;
					*rtime = tstats->ac_etime * 1e-6;
					return 0;
				}
				len2 += NLA_ALIGN(nattr->nla_len);
				nattr = (struct nlattr *)((char *)nattr + len2);
			}
		}
		nattr = (struct nlattr *)(GENLMSG_DATA(&msg) + len);
	}
	return 0;
}

int send_nl_cmd(int sd, __u16 nlmsg_type, __u32 nlmsg_pid, __u8 genl_cmd, __u16 nla_type, 
				void *nla_data, int nla_len) {
	struct nlattr *na;
	struct sockaddr_nl nladdr;
	int r, buflen;
	char *buf;
	struct nl_msg_template msg;

	msg.n.nlmsg_len = NLMSG_LENGTH(GENL_HDRLEN);
	msg.n.nlmsg_type = nlmsg_type;
	msg.n.nlmsg_flags = NLM_F_REQUEST;
	msg.n.nlmsg_seq = 0;
	msg.n.nlmsg_pid = nlmsg_pid;
	msg.g.cmd = genl_cmd;
	msg.g.version = 0x1;
	na = (struct nlattr *) GENLMSG_DATA(&msg);
	na->nla_type = nla_type;
	na->nla_len = nla_len + 1 + NLA_HDRLEN;
	memcpy(NLA_DATA(na), nla_data, nla_len);
	msg.n.nlmsg_len += NLMSG_ALIGN(na->nla_len);

	buf = (char *) &msg;
	buflen = msg.n.nlmsg_len ;
	memset(&nladdr, 0, sizeof(nladdr));
	nladdr.nl_family = AF_NETLINK;
	while ((r = sendto(sd, buf, buflen, 0, (struct sockaddr *) &nladdr,
					   sizeof(nladdr))) < buflen) {
		if (r > 0) {
			buf += r;
			buflen -= r;
		} else if (errno != EAGAIN)
			return -1;
	}
	return 0;
}
#endif
